/**
 * @file acl_resource_manager.h
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#ifndef ASCEND_ACL_RESOURCE_MANAGER_H
#define ASCEND_ACL_RESOURCE_MANAGER_H

#include <map>
#include <mutex>
#include <memory>
#include <string>
#include <vector>
#include <unordered_map>
#include <atomic>
#include "acl/acl.h"
#include "acl/acl_base.h"
#include "graph/ge_attr_value.h"
#include "mmpa/mmpa_api.h"
#include "types/op_model.h"
#include "types/acl_op.h"
#include "single_op/op_model_cache.h"
#include "single_op/compile/op_compiler.h"
#include "utils/acl_op_map.h"
#include "utils/string_utils.h"
#include "framework/runtime/mem_allocator.h"
#include "framework/runtime/model_v2_executor.h"
#include "framework/runtime/stream_executor.h"
#include "external/ge/ge_allocator.h"

namespace acl {
struct BundleInfo {
    uint32_t modelId = 0;
};

class ACL_FUNC_VISIBILITY AclResourceManager {
public:
    ~AclResourceManager();

    static AclResourceManager &GetInstance()
    {
        static AclResourceManager instance;
        return instance;
    }

    // executor
    bool IsRuntimeV2Enable(bool isModel) const
    {
        return isModel ? enableRuntimeV2ForModel_ : enableRuntimeV2ForSingleOp_;
    }

    void AddExecutor(uint32_t &modelId, std::unique_ptr<gert::ModelV2Executor> &&executor,
                     const std::shared_ptr<gert::RtSession> &rtSession);
    std::shared_ptr<gert::ModelV2Executor> GetExecutor(const uint32_t modelId);
    aclError DeleteExecutor(const uint32_t modelId);

    void GetRuntimeV2Env();
    std::shared_ptr<gert::RtSession> CreateRtSession();
    std::shared_ptr<gert::RtSession> GetRtSession(const uint32_t rtSessionId);
    // allocator
    gert::Allocators *GetAllocators(const aclrtStream stream, bool createDefaultAllocator = true);
    void CleanAllocators(const void * const cacheKey);

    static void *GetKeyByStreamOrDefaultStream(const aclrtStream stream);

    // op_model
    aclError LoadAllModels(const std::string &modelDir);

    aclError LoadModelFromMem(const void *const model,
                              const size_t modelSize,
                              const bool isStatic = false);

    aclError LoadModelFromSharedMem(const std::shared_ptr<void> &model,
                                    const size_t modelSize,
                                    const AclOp *const aclOp,
                                    const bool isStatic = false);

    aclError GetOpModel(AclOp &aclOp);

    aclError MatchOpModel(const AclOp &aclOp, OpModel &opModel, bool &isDynamic);

    aclError UpdateRT2Executor(const uint64_t id, const std::shared_ptr<gert::StreamExecutor> &executor);

    std::shared_ptr<std::mutex> GetCacheMutex(const uint64_t id);

    std::shared_ptr<gert::StreamExecutor> GetRT2Executor(const uint64_t id);

    aclError UnloadModelData(const uint64_t id);

    aclError CleanRT2Executor(rtStream_t stream);

    aclError HandleMaxOpQueueConfig(const char_t *const configPath);

    aclError SetHostMemToConst(const AclOp &aclopHostMemToConst, bool &isExistConst) const;

    static aclError SetTensorConst(aclTensorDesc *const desc, const aclDataBuffer *const dataBuffer);

    static void SetCompileFlag(const int32_t flag);

    static void SetJitCompileFlag(const int32_t flag);

    void SetMaxOpNum(const uint64_t maxOpNum);

    aclError CreateExternalAllocator(const aclrtStream stream, void *allocatorDesc);

    aclError DeleteExternalAllocator(const aclrtStream stream);

    std::shared_ptr<ge::Allocator> GetDeviceAllocator(const aclrtStream stream, bool createDefaultAllocator = true);

    void AddBundleInfo(const uint32_t bundleId, const std::vector<BundleInfo> &bundleInfos);

    aclError GetBundleInfo(const uint32_t bundleId, std::vector<BundleInfo> &bundleInfos);

    void DeleteBundleInfo(const uint32_t bundleId);

    bool IsBundleInnerId(const uint32_t modelId);

    void HandleReleaseSourceByStream(aclrtStream stream, bool isCreate);
    void HandleReleaseSourceByDevice(uint32_t devId, bool isReset) const;

    aclError CreateRT2Executor(std::shared_ptr<gert::StreamExecutor> &streamExecutor, rtStream_t stream,
                               const gert::ModelExecuteArg &arg, gert::ModelV2Executor *&executor);

private:
    AclResourceManager() = default;

    // op_model
    using ModelMap = AclOpMap<std::shared_ptr<OpModelDef>>;

    aclError RegisterModel(OpModelDef &&modelConfig,
                           ModelMap &opModelDefs,
                           const bool isDynamic,
                           bool &isRegistered,
                           const bool isStaticRegister = true);

    aclError ReadModelDefs(const std::string &configPath,
                           std::vector<OpModelDef> &configList);

    aclError BuildOpModel(const AclOp &aclOp);

    static bool OmFileFilterFn(const std::string &fileName);

    static bool IsDynamicOpModel(const OpModelDef &modelDef);

    static bool IsDynamicOpModel(const AclOp &aclOp);

    aclError MatchStaticOpModel(const AclOp &aclOp, OpModel &opModel, bool &isDynamic, bool &isNeedMatchDymaic);
    aclError MatchDynamicOpModel(const AclOp &aclOp, OpModel &opModel, bool &isDynamic);

private:
    // executor
    // model id 0 is invalid value
    std::unordered_map<uint32_t, std::shared_ptr<gert::ModelV2Executor>> executorMap_{{0U, nullptr}};
    std::atomic_uint32_t modelIdGenerator_ {std::numeric_limits<uint32_t>::max() / 2U};
    std::atomic_uint64_t sessionIdGenerator_ {std::numeric_limits<uint64_t>::max() / 2U};
    bool enableRuntimeV2ForModel_ = true;
    bool enableRuntimeV2ForSingleOp_ = true;
    std::mutex mutex_;
    std::map<uint32_t, std::shared_ptr<gert::RtSession>> rtSessionMap_;

    // allocator
    // note: op_model的executor的释放依赖allocator
    std::map<const void *, std::shared_ptr<gert::Allocators> > streamAllocator_; // cacheKey + allocator_handle
    std::recursive_mutex streamAllocatorMutex_;
    std::map<const void *, std::shared_ptr<ge::Allocator> > streamExternalAllocator_; // cacheKey + allocator_handle
    std::mutex streamExternalAllocatorMutex_;

    // op_model
    ModelMap opModels_;
    OpModelCache modelCache_;

    std::unordered_map<uint32_t, std::vector<BundleInfo>> bundleInfos_;
    std::unordered_set<uint32_t> bundleInnerIds_;
};
}

#endif // ASCEND_ACL_RESOURCE_MANAGER_H
